package com.sunop.panticommon.annotation.Aspect;

import cn.dev33.satoken.stp.StpUtil;
import com.github.benmanes.caffeine.cache.Cache;
import com.sunop.panticommon.annotation.RateLimit;
import jakarta.servlet.http.HttpServletRequest;
import lombok.RequiredArgsConstructor;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.util.concurrent.atomic.AtomicInteger;

@Aspect
@Component
@RequiredArgsConstructor
public class RateLimitAspect {
    private final Cache<String, AtomicInteger> userQpsCache;
    private final Cache<String, AtomicInteger> ipRequestCache;
    private final Cache<String, Long> ipBlacklistCache;

    @Around("@annotation(rateLimit)")
    public Object doRateLimit(ProceedingJoinPoint pjp, RateLimit rateLimit) throws Throwable {
        // 1. 获取当前用户ID（Sa-Token）
        String userId = StpUtil.getLoginIdAsString();
        
        // 2. 获取客户端IP
        String ip = getClientIp();

        // 3. IP 黑名单检查
        if (isIpBanned(ip, rateLimit.ipBanMinutes())) {
            throw new RuntimeException("IP已被封禁，请稍后再试！");
        }

        // 4. 用户 QPS 限流
        if (isUserOverQps(userId, rateLimit.userQps())) {
            StpUtil.kickout(userId); // 强制下线
            throw new RuntimeException("操作过于频繁，账号已被强制下线！");
        }

        // 5. IP 请求计数检查
        if (isIpOverLimit(ip, rateLimit.ipMaxRequests())) {
            ipBlacklistCache.put(ip, System.currentTimeMillis()); // 加入黑名单
            throw new RuntimeException("请求过于频繁，IP已被封禁！");
        }

        // 6. 放行请求
        return pjp.proceed();
    }

    // 获取客户端IP（需适配你的框架）
    private String getClientIp() {
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        return request.getRemoteAddr();
    }

    // 检查IP是否在黑名单
    private boolean isIpBanned(String ip, int banMinutes) {
        Long banTime = ipBlacklistCache.getIfPresent(ip);
        if (banTime != null) {
            long currentTime = System.currentTimeMillis();
            long banDuration = banMinutes * 60 * 1000L;
            return (currentTime - banTime) < banDuration;
        }
        return false;
    }

    // 检查用户QPS是否超限
    private boolean isUserOverQps(String userId, int maxQps) {
        AtomicInteger count = userQpsCache.get(userId, k -> new AtomicInteger(0));
        return count.incrementAndGet() > maxQps;
    }

    // 检查IP请求是否超限
    private boolean isIpOverLimit(String ip, int maxRequests) {
        AtomicInteger count = ipRequestCache.get(ip, k -> new AtomicInteger(0));
        return count.incrementAndGet() > maxRequests;
    }
}